Skip to content

Fix manual vLLM Qwen3 sharding bug when trainer export the weights #229

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Merged
merged 12 commits into from
Sep 26, 2025
Merged

Fix manual vLLM Qwen3 sharding bug when trainer export the weights #229

merged 12 commits into from
Sep 26, 2025

Conversation

JenniferWang
Copy link
Contributor

@JenniferWang JenniferWang commented Sep 24, 2025 •
edited
Loading

Thanks to @casteryh 's PR of fixing the integration test #217, we are able to see the buggy behavior of the manual sharding logic.

This diff basically showcases why the test failed and how easy it is to introduce silent data correctness issue if we keep pursuing this route of manual sharding.

The plan is land these two PRs ASAP

  1. Fix the broken test: integration test for weight sync that actually tests behavior #217
  2. Replace the manual sharding with vanilla load weight calls: Weight loading working correctly with tp: use vllm builtin load_weights() #184

At the same time, this PR fixes the existing Qwen3 (non MoE) sharding.

Before the fix:

  • test_policy_update integration test fails when tp > 1
  • grpo (Qwen 1.7B) loss function super high at the start of the trainer
image https://meta.wandb.io/jiyue/grpo-training/runs/hh3bht2w?nw=nwuserjiyue

After the fix:

  • test_policy_update integration test passes when tp > 1
  • grpo (Qwen 1.7B) loss function is much more reasonable??
image https://meta.wandb.io/jiyue/grpo-training/runs/gnvlw7dc?nw=nwuserjiyue

@meta-cla meta-cla bot added the CLA Signed This label is managed by the Meta Open Source bot. label Sep 24, 2025
@JenniferWang JenniferWang linked an issue Sep 24, 2025 that may be closed by this pull request
Give a better solution to manually calculated vLLM sharding logic in sharding.py #174
Closed
@allenwang28
Copy link
Contributor

@JenniferWang life pro tip, you should be able to do pre-commit install (assuming you've pip install pre-commit) which will handle the linter for you automatically every time you do a git commit

@JenniferWang JenniferWang changed the title Demonstrate the vLLM sharding bug when trainer export the weights Fix manual vLLM Qwen3 sharding bug when trainer export the weights Sep 26, 2025
@JenniferWang JenniferWang merged commit 825596b into meta-pytorch:main Sep 26, 2025
5 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@joecummings joecummings joecummings approved these changes

Assignees

No one assigned

Labels

CLA Signed This label is managed by the Meta Open Source bot.

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

Give a better solution to manually calculated vLLM sharding logic in sharding.py

3 participants

@JenniferWang @allenwang28 @joecummings